from PIL import Image
from torch.utils.data import Dataset
import random
import PIL.ImageOps
import numpy as np
import torch
class SiameseNetworkDataset(Dataset):
    def __init__(self, imageFolderDataset,transform=None,should_invert=True):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert
    def __getitem__(self, index):
        img0_tuple = random.choice( self.imageFolderDataset.imgs)
        # 需要确保大约50%的图像属于同一类
        should_get_same_class = random.randint(0, 1)
        if should_get_same_class:
            while True:
                # 保持循环直到找到相同的类图像
                img1_tuple = random.choice(
                              self.imageFolderDataset.imgs)
                if img0_tuple[1] == img1_tuple[1]:
                    break
        else:
            while True:
                # 保持循环直到找到不同的类图像
                img1_tuple=random.choice( self.imageFolderDataset.imgs)
                if img0_tuple[1] != img1_tuple[1]:
                    break
        img0 = Image.open(img0_tuple[0])
        img1 = Image.open(img1_tuple[0])
        img0 = img0.convert("RGB")
        img1 = img1.convert("RGB")
        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)
        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)
        return img0,img1,torch.from_numpy(
              np.array([int(img1_tuple[1]==img0_tuple[1])],dtype=np.float32))
    def __len__(self):
        return len(self.imageFolderDataset.imgs)
